#pragma once

#include <vector>
#include "helper/helper.h"
#include <boost/mpi.hpp>
extern boost::mpi::communicator* p_world;
extern po::variables_map vm;
#include "mpi_helper.h"

/**
  * Indexing neighbors for each node
  * depends on graph type (grid or equal)
  */
class GraphHelper
{
  enum GraphType { GRID, EQUAL } m_Graph;
  int m_GridSize, m_Row, m_Col;
  int m_Split, m_Degree;
public:
  GraphHelper() 
  {
    PO_VARIABLE(graph, vm, "graph", std::string);
    if (graph.compare("grid") == 0) {
      m_Graph = GRID;
      m_GridSize = (int) floor(sqrt(MPI_SIZE));
      id2rc(m_GridSize, MPI_RANK, m_Row, m_Col); // index of current node is MPI_RANK
                                                 // total number of nodes is MPI_SIZE
      MSG_VERBOSE(vm, "Grid " << m_GridSize << " " << m_Row << " " << m_Col);
      if (m_GridSize * m_GridSize != MPI_SIZE) 
        FATAL_MASTER(0, "Number of nodes is not a square number");    
    }
    else if (graph.compare("equal") == 0) {
      m_Graph = EQUAL;
      m_Split = MPI_SIZE / 2;
      m_Degree = vm["degree"].as<int>();
    }
    else FATAL_MASTER(0, "unknown graph type");
  }
  
  void BuildGraph(std::vector<int>& neighbors)
  {
    neighbors.clear();
    switch (m_Graph) 
    {
      case GRID:
        BuildGrid(neighbors);
        break;
      case EQUAL:
        BuildEqualDegree(neighbors);
        break;
      default:
        FATAL_MASTER(0, "unknown graph type");
    }
  }
  
  bool IsLeftSide()
  {
    switch (m_Graph) 
    {
      case GRID:
          return abs(m_Row-m_Col) % 2 == 0;
      case EQUAL:
        return MPI_RANK < m_Split;
      default:
        return false;
    }
  }
  
  static int rc2id(int n, int r, int c) {
    return r*n+c;
  }
  
  static void id2rc(int n, int id, int& r, int& c) {
    r = id / n;
    c = id % n;
  }
  
  void BuildGrid(std::vector<int>& neighbors)
  {
    if (m_Row > 0) neighbors.push_back(rc2id(m_GridSize, m_Row-1,m_Col));
    if (m_Row < m_GridSize-1) neighbors.push_back(rc2id(m_GridSize, m_Row+1,m_Col));
    if (m_Col > 0) neighbors.push_back(rc2id(m_GridSize, m_Row,m_Col-1));
    if (m_Col < m_GridSize-1) neighbors.push_back(rc2id(m_GridSize, m_Row,m_Col+1));
  }
  
  void BuildEqualDegree(std::vector<int>& neighbors)
  {
    int d2 = m_Degree / 2;
    if (IsLeftSide()) {
      int start = MPI_RANK+m_Split-d2;
      for (int i = 0; i < m_Degree; i++, start++) 
        if (start >= m_Split && start < MPI_SIZE) neighbors.push_back(start);
    }
    else
    {
      int end = MPI_RANK-m_Split+d2;
      for (int i = 0; i < m_Degree; i++, end--)
        if (end >= 0 && end < m_Split) neighbors.push_back(end);
    }
  }
};

